# apps/orchestrator/matpo_orchestrator.py
import asyncio
import aiohttp
import numpy as np
from typing import Optional, Dict, Any, List, Tuple
from datetime import datetime
import threading
import random
from dataclasses import dataclass
from loguru import logger
import json

# 导入项目核心模块和常量
from utils.logger import Logger
from utils.exceptions import (
    RedisError, LockAcquireError, ParamError,
    AgentNotFoundError, AgentCallError
)
from utils.constants import (
    AgentType, TaskState, ErrorCode,
    AGENT_STATE_UPDATE_CHANNEL, TASK_STATE_UPDATE_CHANNEL,
    TaskStatusConst, TaskPriority
)
from middleware.redis_adapter import get_redis_adapter
from states.agent_state import AgentState
from states.task_state import TaskStateManager
from states.context_partition_manager import context_partition_manager, ContextPartitionType
from app.core.db_utils import DatabaseUtils
from apps.orchestrator.task_review_manager import task_review_manager
from middleware.model_interface import get_model_client
from middleware.model_strategy import model_selector, TaskType
from config.model_config import ModelBackend, model_config_manager
from communication.agent_prune import agent_prune_manager


@dataclass
class AgentCapability:
    """智能体能力描述"""
    agent_type: str
    instance_id: str
    capability_vector: List[float]  # 能力向量，表示在不同领域的能力强度
    success_rate: float  # 历史成功率
    response_time: float  # 平均响应时间
    current_load: int  # 当前负载
    last_heartbeat: int  # 最后心跳时间


class MATPOOrchestratorAgent:
    """基于MATPO训练方法的动态角色分配协调器"""

    def __init__(self, redis_client=None):
        # 依赖组件
        self.redis_client = redis_client or get_redis_adapter()
        self.logger = Logger.get_logger(agent_name="matpo_orchestrator_agent")

        # 缓存：健康Agent实例和能力信息
        self.healthy_agents: Dict[str, List[Dict[str, Any]]] = {}  # {agent_type: [instance_info]}
        self.agent_capabilities: Dict[str, AgentCapability] = {}  # {instance_id: AgentCapability}
        self.capability_dimension = 6  # 能力向量维度
        self.cache_lock = threading.Lock()  # 缓存更新锁（线程安全）

        # MATPO相关参数
        self.epsilon = 0.1  # 探索率
        self.learning_rate = 0.01  # 学习率
        self.discount_factor = 0.9  # 折扣因子
        self.task_type_embeddings = self._init_task_type_embeddings()  # 预定义任务类型嵌入

        # 异步HTTP会话（连接池复用）
        self._session: Optional[aiohttp.ClientSession] = None
        
        # 启动Agent状态订阅（实时感知实例上下线）
        self._start_agent_state_subscribe()
        self.logger.success("MATPO协调器初始化完成，已启动Agent状态订阅")

    def _init_task_type_embeddings(self) -> Dict[str, List[float]]:
        """初始化任务类型嵌入向量"""
        return {
            "text_qa": [0.9, 0.1, 0.2, 0.3, 0.1, 0.4],
            "doc_qa": [0.8, 0.7, 0.3, 0.2, 0.1, 0.5],
            "tool_call": [0.3, 0.2, 0.9, 0.8, 0.7, 0.3],
            "multi_modal": [0.5, 0.8, 0.7, 0.6, 0.9, 0.4],
            "doc_segment": [0.6, 0.9, 0.2, 0.1, 0.3, 0.6],
        }

    async def __aenter__(self):
        """异步上下文管理器进入"""
        if not self._session or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=30),
                connector=aiohttp.TCPConnector(
                    limit=100,  # 连接池大小
                    limit_per_host=20,  # 每个主机最大连接数
                    ttl_dns_cache=300,  # DNS缓存时间
                    use_dns_cache=True,
                ),
                headers={
                    'User-Agent': 'MAESS-MATPOOrchestrator/1.0',
                    'Content-Type': 'application/json',
                }
            )
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器退出"""
        if self._session and not self._session.closed:
            await self._session.close()

    @property
    async def session(self) -> aiohttp.ClientSession:
        """获取或创建HTTP会话"""
        if not self._session or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=30),
                connector=aiohttp.TCPConnector(
                    limit=100,
                    limit_per_host=20,
                    ttl_dns_cache=300,
                    use_dns_cache=True,
                ),
                headers={
                    'User-Agent': 'MAESS-MATPOOrchestrator/1.0',
                    'Content-Type': 'application/json',
                }
            )
        return self._session

    def _start_agent_state_subscribe(self) -> None:
        """启动Agent状态订阅（实时更新健康实例缓存）"""
        stop_event = threading.Event()

        def agent_state_callback(data: Dict[str, Any]) -> None:
            """Agent状态更新回调（实例注册/注销/更新时触发）"""
            agent_type = data.get("agent_type", "")
            instance_id = data.get("instance_id", "")
            action = data.get("action", "update")
            status = data.get("status", "")

            self.logger.debug(
                f"收到Agent状态更新: action={action} | type={agent_type} | instance={instance_id} | status={status}")

            # 加锁更新健康实例缓存
            with self.cache_lock:
                if action in ["register", "update"] and status == "healthy":
                    # 新增/更新健康实例
                    if agent_type not in self.healthy_agents:
                        self.healthy_agents[agent_type] = []
                    
                    # 移除旧实例信息（避免重复）
                    self.healthy_agents[agent_type] = [
                        ins for ins in self.healthy_agents[agent_type]
                        if ins["instance_id"] != instance_id
                    ]
                    
                    # 添加新实例信息
                    instance_info = {
                        "instance_id": instance_id,
                        "host_ip": data.get("host_ip", ""),
                        "port": data.get("port", 0),
                        "last_heartbeat": data.get("last_heartbeat", 0)
                    }
                    self.healthy_agents[agent_type].append(instance_info)
                    
                    # 更新或创建Agent能力信息
                    current_tasks = int(data.get("current_tasks", 0))
                    success_rate = float(data.get("success_rate", 0.5))
                    response_time = float(data.get("avg_response_time", 1.0))
                    
                    # 初始化或更新能力向量
                    if instance_id in self.agent_capabilities:
                        # 已有记录，使用现有能力向量
                        capability_vector = self.agent_capabilities[instance_id].capability_vector
                    else:
                        # 新实例，基于agent_type初始化能力向量
                        base_capability = self._init_task_type_embeddings().get(agent_type, [0.5] * self.capability_dimension)
                        # 添加一些随机扰动
                        capability_vector = [min(1.0, max(0.0, c + random.uniform(-0.1, 0.1))) for c in base_capability]
                    
                    self.agent_capabilities[instance_id] = AgentCapability(
                        agent_type=agent_type,
                        instance_id=instance_id,
                        capability_vector=capability_vector,
                        success_rate=success_rate,
                        response_time=response_time,
                        current_load=current_tasks,
                        last_heartbeat=data.get("last_heartbeat", 0)
                    )
                    
                elif action == "unregister" or status != "healthy":
                    # 移除注销/不健康实例
                    if agent_type in self.healthy_agents:
                        self.healthy_agents[agent_type] = [
                            ins for ins in self.healthy_agents[agent_type]
                            if ins["instance_id"] != instance_id
                        ]
                        # 空列表时删除键
                        if not self.healthy_agents[agent_type]:
                            del self.healthy_agents[agent_type]
                    
                    # 移除能力信息
                    if instance_id in self.agent_capabilities:
                        del self.agent_capabilities[instance_id]

        # 启动订阅线程（守护线程，随主进程退出）
        subscribe_thread = threading.Thread(
            target=self.redis_client.subscribe,
            args=(AGENT_STATE_UPDATE_CHANNEL, agent_state_callback, stop_event),
            daemon=True
        )
        subscribe_thread.start()
        self.logger.info(f"Agent状态订阅线程已启动，频道={AGENT_STATE_UPDATE_CHANNEL}")

    def _calculate_task_embedding(self, task_input: Dict[str, Any]) -> List[float]:
        """计算任务嵌入向量"""
        task_type = task_input.get("task_type", "single")
        agent_type = task_input.get("agent_type", AgentType.TEXT_QA.value)
        input_text = task_input.get("input", "")
        
        # 基于预定义的任务类型嵌入
        base_embedding = self._init_task_type_embeddings().get(agent_type, [0.5] * self.capability_dimension)
        
        # 如果有输入文本，尝试生成语义嵌入并融合
        if input_text:
            try:
                # 生成文本嵌入
                text_embedding = DatabaseUtils.generate_embedding(input_text)
                # 降维或截断到capability_dimension维度
                if len(text_embedding) > self.capability_dimension:
                    # 使用主成分分析简单降维
                    text_embedding = text_embedding[:self.capability_dimension]
                elif len(text_embedding) < self.capability_dimension:
                    # 填充0
                    text_embedding.extend([0.0] * (self.capability_dimension - len(text_embedding)))
                
                # 归一化
                text_embedding = [min(1.0, max(0.0, x/20)) for x in text_embedding]  # 简单归一化
                
                # 融合任务类型嵌入和文本语义嵌入
                final_embedding = [0.7 * base + 0.3 * text for base, text in zip(base_embedding, text_embedding)]
                return final_embedding
            except Exception as e:
                self.logger.warning(f"生成任务嵌入失败: {str(e)}")
        
        return base_embedding

    def _calculate_capability_score(self, agent_capability: AgentCapability, task_embedding: List[float]) -> float:
        """计算智能体对任务的能力匹配得分"""
        # 计算能力向量和任务嵌入的余弦相似度
        if len(agent_capability.capability_vector) != len(task_embedding):
            return 0.0
        
        # 余弦相似度计算
        dot_product = sum(a * b for a, b in zip(agent_capability.capability_vector, task_embedding))
        norm_a = np.linalg.norm(agent_capability.capability_vector)
        norm_b = np.linalg.norm(task_embedding)
        
        if norm_a == 0 or norm_b == 0:
            similarity = 0.0
        else:
            similarity = dot_product / (norm_a * norm_b)
        
        # 考虑其他因素：成功率、响应时间、负载
        # 成功率权重0.2，响应时间权重0.1（越小越好），负载权重0.1（越小越好）
        success_rate_factor = agent_capability.success_rate * 0.2
        response_time_factor = max(0.1, 1.0 - agent_capability.response_time / 10.0) * 0.1
        load_factor = max(0.1, 1.0 - agent_capability.current_load / 10.0) * 0.1
        
        # 综合得分
        total_score = 0.6 * similarity + success_rate_factor + response_time_factor + load_factor
        
        return total_score

    def _select_agent_by_matpo(self, task_embedding: List[float], agent_type: str) -> Optional[Dict[str, Any]]:
        """基于MATPO算法选择最优Agent实例"""
        # 1. 获取该类型的健康实例
        with self.cache_lock:
            healthy_instances = self.healthy_agents.get(agent_type, [])
        
        if not healthy_instances:
            self.logger.warning(f"无{agent_type}类型的健康实例")
            return None
        
        # 2. 计算每个实例的能力得分
        scored_instances = []
        for instance in healthy_instances:
            instance_id = instance["instance_id"]
            if instance_id in self.agent_capabilities:
                capability = self.agent_capabilities[instance_id]
                score = self._calculate_capability_score(capability, task_embedding)
                scored_instances.append((instance, score))
        
        if not scored_instances:
            self.logger.warning(f"无{agent_type}类型的能力信息")
            # 如果没有能力信息，回退到负载均衡
            return self._select_best_agent_by_load(agent_type)
        
        # 3. MATPO策略：ε-贪婪算法
        if random.random() < self.epsilon:
            # 探索：随机选择一个实例
            self.logger.debug("MATPO策略：探索模式")
            selected_instance, _ = random.choice(scored_instances)
        else:
            # 利用：选择得分最高的实例
            self.logger.debug("MATPO策略：利用模式")
            selected_instance, _ = max(scored_instances, key=lambda x: x[1])
        
        self.logger.debug(f"MATPO选择实例: {selected_instance['instance_id']} 类型: {agent_type}")
        return selected_instance

    def _select_best_agent_by_load(self, agent_type: str) -> Optional[Dict[str, Any]]:
        """基于负载均衡选择最优Agent实例（回退策略）"""
        with self.cache_lock:
            healthy_instances = self.healthy_agents.get(agent_type, [])
        
        if not healthy_instances:
            return None
        
        # 选择当前任务数最少的实例
        best_instance = None
        min_task_count = float("inf")
        for instance in healthy_instances:
            instance_id = instance["instance_id"]
            redis_key = f"agent:state:{agent_type}:{instance_id}"
            # 从Redis获取当前任务数
            current_tasks = self.redis_client.hget(redis_key, "current_tasks")
            current_tasks = int(current_tasks) if current_tasks else 0

            # 选择任务数最少的实例
            if current_tasks < min_task_count:
                min_task_count = current_tasks
                best_instance = instance
                best_instance["current_tasks"] = current_tasks
        
        return best_instance

    def _update_agent_capabilities(self, instance_id: str, success: bool, task_embedding: List[float], reward: float) -> None:
        """根据任务执行结果更新智能体能力（MATPO学习）"""
        if instance_id not in self.agent_capabilities:
            return
        
        agent_cap = self.agent_capabilities[instance_id]
        
        # 1. 更新成功率和响应时间
        if success:
            # 简单的滑动平均更新
            agent_cap.success_rate = 0.9 * agent_cap.success_rate + 0.1 * 1.0
        else:
            agent_cap.success_rate = 0.9 * agent_cap.success_rate + 0.1 * 0.0
        
        # 2. 根据MATPO算法更新能力向量
        # 基于任务嵌入和奖励调整能力向量
        for i in range(len(agent_cap.capability_vector)):
            # 能力向量更新公式：cap[i] = cap[i] + α * reward * task_embedding[i]
            agent_cap.capability_vector[i] += self.learning_rate * reward * task_embedding[i]
            # 确保在有效范围内
            agent_cap.capability_vector[i] = min(1.0, max(0.0, agent_cap.capability_vector[i]))
        
        self.logger.debug(f"更新实例{instance_id}能力: 成功率={agent_cap.success_rate:.2f}, 能力向量={agent_cap.capability_vector[:3]}...")

    async def _call_agent_instance_async(self, instance: Dict[str, Any], task_params: Dict[str, Any]) -> Dict[str, Any]:
        """异步调用Agent实例执行任务"""
        agent_type = task_params.get("agent_type", "")
        instance_id = instance["instance_id"]
        host_ip = instance["host_ip"]
        port = instance["port"]

        self.logger.info(f"异步调用Agent实例: {instance_id} | 参数: {task_params}")

        # 构建异步HTTP请求URL和参数
        url = f"http://{host_ip}:{port}/api/run"
        request_data = {
            "task_id": task_params.get("task_id", ""),
            "agent_type": agent_type,
            "input": task_params.get("input", ""),
            "timestamp": int(datetime.now().timestamp())
        }

        try:
            session = await self.session
            async with session.post(url, json=request_data) as response:
                # 设置超时时间
                response_data = await asyncio.wait_for(response.json(), timeout=30.0)
                
                if response.status == 200:
                    return {
                        "success": True,
                        "data": response_data,
                        "error_code": 0,
                        "error_msg": ""
                    }
                else:
                    return {
                        "success": False,
                        "data": None,
                        "error_code": response.status,
                        "error_msg": f"HTTP {response.status}: {response_data.get('error', 'Unknown error')}"
                    }

        except asyncio.TimeoutError:
            self.logger.error(f"调用Agent实例{instance_id}超时")
            return {
                "success": False,
                "data": None,
                "error_code": ErrorCode.AGENT_CALL_ERROR.value,
                "error_msg": "Agent调用超时(30s)"
            }
        except aiohttp.ClientError as e:
            self.logger.error(f"调用Agent实例{instance_id}网络错误: {str(e)}")
            return {
                "success": False,
                "data": None,
                "error_code": ErrorCode.AGENT_CALL_ERROR.value,
                "error_msg": f"网络错误: {str(e)}"
            }
        except Exception as e:
            self.logger.error(f"调用Agent实例{instance_id}失败: {str(e)}")
            return {
                "success": False,
                "data": None,
                "error_code": ErrorCode.AGENT_CALL_ERROR.value,
                "error_msg": f"调用失败: {str(e)}"
            }

    async def _retry_task_async(self, task_id: str, agent_type: str, task_params: Dict[str, Any], task_embedding: List[float]) -> Tuple[bool, Dict[str, Any]]:
        """异步任务重试逻辑"""
        task_manager = TaskStateManager(task_id=task_id, redis_client=self.redis_client)
        retry_count = task_manager.get_state()["retry_count"]
        max_retries = task_manager.get_state()["max_retries"] or 3

        if retry_count >= max_retries:
            self.logger.error(f"任务{task_id}已达到最大重试次数({max_retries})，停止重试")
            task_manager.update_state(
                state=TaskState.FAILED.value,
                error_code=ErrorCode.TASK_EXECUTE_ERROR.value,
                error_msg=f"Max retries ({max_retries}) exceeded"
            )
            return False, {"success": False, "error_msg": "Max retries exceeded"}

        # 重试计数+1，更新任务状态为重试中
        task_manager.update_state(state=TaskState.RETRYING.value)
        self.logger.info(f"任务{task_id}开始第{retry_count + 1}次重试")

        try:
            # 重新选择Agent实例（避免复用失败的实例）
            instance = self._select_agent_by_matpo(task_embedding, agent_type)
            if not instance:
                raise AgentNotFoundError(
                    message=f"No available agent instance for retry",
                    context={"agent_type": agent_type}
                )
            
            # 重新调用
            result = await self._call_agent_instance_async(instance, task_params)
            if result["success"]:
                # 重试成功，更新任务状态
                task_manager.update_state(
                    state=TaskState.SUCCESS.value,
                    output_result=result["data"],
                    agent_instance_id=instance["instance_id"]
                )
                
                # 计算奖励并更新能力
                reward = 1.0  # 成功奖励
                self._update_agent_capabilities(instance["instance_id"], True, task_embedding, reward)
                
                return True, result
            else:
                # 重试失败，计算惩罚并更新能力
                reward = -0.5  # 失败惩罚
                if instance and "instance_id" in instance:
                    self._update_agent_capabilities(instance["instance_id"], False, task_embedding, reward)
                
                return await self._retry_task_async(task_id, agent_type, task_params, task_embedding)
        except Exception as e:
            self.logger.error(f"任务{task_id}第{retry_count + 1}次重试失败: {str(e)}")
            return await self._retry_task_async(task_id, agent_type, task_params, task_embedding)

    def _parse_task(self, task_input: Dict[str, Any]) -> List[Dict[str, Any]]:
        """解析任务（将复杂任务拆分为多个子任务，支持单/多Agent协作）"""
        task_type = task_input.get("task_type", "single")  # single/multi
        input_data = task_input.get("input", "")
        task_id = task_input.get("task_id", "")

        if task_type == "single":
            # 单Agent任务：直接指定Agent类型
            agent_type = task_input.get("agent_type", AgentType.TEXT_QA.value)
            return [{
                "sub_task_id": f"sub_{task_id}_001",
                "agent_type": agent_type,
                "input": input_data,
                "task_id": task_id
            }]
        elif task_type == "multi":
            # 多Agent协作任务：根据任务特性动态拆分
            # 这里可以基于MATPO的角色分配逻辑进行更智能的拆分
            # 目前实现一个简单的两阶段协作模式：信息获取→综合分析
            return [
                {
                    "sub_task_id": f"sub_{task_id}_001",
                    "agent_type": AgentType.TOOL_CALL.value,
                    "input": f"获取{input_data}相关信息",
                    "task_id": task_id
                },
                {
                    "sub_task_id": f"sub_{task_id}_002",
                    "agent_type": AgentType.TEXT_QA.value,
                    "input": input_data,
                    "task_id": task_id
                }
            ]
        else:
            raise ParamError(
                message="Invalid task_type (only support 'single'/'multi')",
                context={"task_type": task_type, "error_code": ErrorCode.TASK_PARAM_ERROR.value}
            )

    async def execute_task_async(
        self, 
        task_input: Dict[str, Any],
        force_backend: Optional[ModelBackend] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """异步执行任务（核心入口：解析→分发→执行→整合→重试）"""
        # 1. 校验任务输入参数
        required_fields = ["task_id", "input"]
        for field in required_fields:
            if field not in task_input:
                raise ParamError(
                    message=f"Task input missing required field: {field}",
                    context={"task_input": task_input, "error_code": ErrorCode.TASK_PARAM_ERROR.value}
                )
        task_id = task_input["task_id"]
        self.logger.info(f"开始异步执行任务: task_id={task_id} | input={task_input['input']}")

        # 任务文本用于模型选择
        task_text = task_input.get("input", "")
        
        # 如果没有指定任务类型，根据内容自动检测
        task_type = None
        if any(keyword in task_text.lower() for keyword in ["代码", "编程", "function", "class"]):
            task_type = TaskType.CODE_GENERATION
        elif any(keyword in task_text.lower() for keyword in ["总结", "摘要", "概括"]):
            task_type = TaskType.SUMMARIZATION
        elif any(keyword in task_text.lower() for keyword in ["图片", "图像", "ocr"]):
            task_type = TaskType.MULTI_MODAL
            
        # 使用模型选择器选择合适的模型
        config_name, selected_backend, client = await model_selector.select_model_for_task(
            task_text=task_text,
            task_type=task_type,
            force_backend=force_backend,
            **kwargs
        )
        
        # 更新AgentPrune的模型后端配置
            try:
                agent_prune_manager.set_model_backend(selected_backend)
                self.logger.debug(f"AgentPrune模型后端设置成功: {selected_backend.value}")
            except Exception as e:
                self.logger.warning(f"AgentPrune模型后端设置失败: {str(e)}")
            
            # 添加上下文到相应分区
            await self._add_context_to_partition(task_id, task_type, task_input, final_result)

        # 2. 初始化任务状态（待执行）
        task_manager = TaskStateManager(task_id=task_id, redis_client=self.redis_client)
        task_manager.update_state(
            state=TaskState.PENDING.value,
            input_params=task_input
        )
        
        # 3. 获取任务相关上下文
        relevant_context = await self._get_relevant_context(task_type, task_input)
        if relevant_context:
            # 将相关上下文添加到任务输入中
            task_input['context_history'] = relevant_context

        try:
            # 3. 计算任务嵌入向量
            task_embedding = self._calculate_task_embedding(task_input)
            self.logger.debug(f"任务嵌入向量: {task_embedding[:3]}...")

            # 4. 解析任务（拆分子任务）
            sub_tasks = self._parse_task(task_input)
            self.logger.debug(
                f"任务解析完成，子任务数: {len(sub_tasks)} | sub_tasks={[t['sub_task_id'] for t in sub_tasks]}")

            # 5. 并发执行所有子任务（使用asyncio.gather实现真正的并行）
            sub_task_coroutines = []
            for sub_task in sub_tasks:
                # 为每个子任务计算嵌入向量
                sub_task_embedding = self._calculate_task_embedding(sub_task)
                sub_task_coroutines.append(self._execute_single_sub_task_async(
                    sub_task, task_id, sub_task_embedding
                ))

            # 并发执行所有子任务
            sub_task_results = await asyncio.gather(*sub_task_coroutines, return_exceptions=True)

            # 处理执行结果
            valid_results = []
            for i, result in enumerate(sub_task_results):
                if isinstance(result, Exception):
                    self.logger.error(f"子任务{sub_tasks[i]['sub_task_id']}执行异常: {str(result)}")
                    # 可以选择重试或标记为失败
                    continue
                elif result and result.get("success"):
                    valid_results.append(result.get("data", {}))

            if not valid_results:
                raise AgentCallError(
                    message="All sub tasks failed",
                    context={"error_msg": "No valid results from any sub task"}
                )

            # 6. 整合所有子任务结果
            final_result = {
                "task_id": task_id,
                "status": "success",
                "sub_task_count": len(sub_tasks),
                "sub_task_results": valid_results,
                "final_answer": self._integrate_results(valid_results, task_input["input"]),
                "executed_at": int(datetime.now().timestamp())
            }
            
            # 7. 调用任务回顾管理器进行结果总结和冗余信息过滤
            try:
                self.logger.info(f"开始任务回顾: task_id={task_id}")
                review_result = await task_review_manager.review_task_results(task_id, model_backend)
                
                # 将回顾结果合并到最终结果中
                final_result["review_summary"] = review_result["summary"]
                final_result["redundancy_rate"] = review_result["redundancy_rate"]
                final_result["filtered_sub_task_ids"] = review_result["filtered_results"]
                
                self.logger.info(f"任务回顾完成: task_id={task_id}, redundancy_rate={review_result['redundancy_rate']}")
            except Exception as e:
                self.logger.warning(f"任务回顾失败: task_id={task_id}, error={str(e)}")
                # 任务回顾失败不应影响主任务结果
                final_result["review_failed"] = True
            
            # 添加模型选择信息
            final_result["model_selection"] = {
                "config_name": config_name,
                "backend": selected_backend.value if selected_backend else "auto",
                "task_type": task_type.value if task_type else "auto_detected"
            }
            
            # 添加AgentPrune统计信息
            try:
                prune_stats = agent_prune_manager.get_pruning_stats()
                final_result["agent_prune_stats"] = prune_stats
            except Exception as e:
                self.logger.warning(f"获取AgentPrune统计信息失败: {str(e)}")

            self.logger.success(f"异步任务{task_id}执行完成 | 最终结果: {final_result['final_answer']}")
            return final_result
        except Exception as e:
            # 7. 任务执行失败：更新状态
            error_msg = str(e)
            self.logger.error(f"异步任务{task_id}执行失败: {error_msg}")
            task_manager.update_state(
                state=TaskState.FAILED.value,
                error_code=ErrorCode.TASK_EXECUTE_ERROR.value,
                error_msg=error_msg
            )
            
            # 记录失败任务的上下文
            await self._add_context_to_partition(task_id, task_type, task_input, {
                'status': 'failed',
                'error_msg': error_msg,
                'task_id': task_id
            })
            return {
                "task_id": task_id,
                "status": "failed",
                "error_msg": error_msg,
                "executed_at": int(datetime.now().timestamp())
            }

    async def _get_relevant_context(self, task_type, task_input: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        获取与当前任务相关的上下文信息
        
        Args:
            task_type: 任务类型
            task_input: 任务输入数据
            
        Returns:
            相关上下文列表
        """
        try:
            # 映射任务类型到分区类型
            partition_map = {
                TaskType.TEXT_QA: ContextPartitionType.TEXT_QA,
                TaskType.CODE_GENERATION: ContextPartitionType.CODE_GENERATION,
                TaskType.DOC_QA: ContextPartitionType.DOC_QA,
                TaskType.MULTIMODAL: ContextPartitionType.MULTIMODAL,
                TaskType.TOOL_CALL: ContextPartitionType.TOOL_CALL,
                TaskType.ORCHESTRATION: ContextPartitionType.ORCHESTRATOR,
                TaskType.SUMMARIZATION: ContextPartitionType.TEXT_QA
            }
            
            partition_type = partition_map.get(task_type, ContextPartitionType.CHAT_HISTORY)
            
            # 获取相关上下文
            relevant_context = await context_partition_manager.get_context_for_task(
                task_type=partition_type,
                limit=10,
                include_shared=True,
                user_id=task_input.get('user_id')
            )
            
            # 格式化上下文
            formatted_context = []
            for ctx in relevant_context:
                if isinstance(ctx.content, dict) and 'result' in ctx.content and 'task_id' in ctx.content:
                    formatted_context.append({
                        'task_id': ctx.content['task_id'],
                        'input': ctx.content.get('input', {}),
                        'result': ctx.content['result'].get('final_answer', '') if isinstance(ctx.content['result'], dict) else '',
                        'timestamp': ctx.timestamp
                    })
            
            self.logger.debug(f"获取到{len(formatted_context)}条相关上下文")
            return formatted_context[:5]  # 最多返回5条
        except Exception as e:
            self.logger.error(f"获取相关上下文失败: {str(e)}")
            return []
            
    async def _add_context_to_partition(self, task_id: str, task_type, 
                                       task_input: Dict[str, Any], result: Dict[str, Any]) -> None:
        """
        将任务上下文添加到适当的分区
        
        Args:
            task_id: 任务ID
            task_type: 任务类型
            task_input: 任务输入数据
            result: 任务结果
        """
        try:
            # 映射任务类型到分区类型
            partition_map = {
                TaskType.TEXT_QA: ContextPartitionType.TEXT_QA,
                TaskType.CODE_GENERATION: ContextPartitionType.CODE_GENERATION,
                TaskType.DOC_QA: ContextPartitionType.DOC_QA,
                TaskType.MULTIMODAL: ContextPartitionType.MULTIMODAL,
                TaskType.TOOL_CALL: ContextPartitionType.TOOL_CALL,
                TaskType.ORCHESTRATION: ContextPartitionType.ORCHESTRATOR,
                TaskType.SUMMARIZATION: ContextPartitionType.TEXT_QA
            }
            
            partition_type = partition_map.get(task_type, ContextPartitionType.CHAT_HISTORY)
            
            # 准备上下文内容
            context_content = {
                'task_id': task_id,
                'task_type': task_type.value if hasattr(task_type, 'value') else str(task_type),
                'input': task_input.get('input', {}),
                'result': result,
                'timestamp': asyncio.get_event_loop().time()
            }
            
            # 确定重要性评分
            importance = 0.5  # 默认重要性
            if result.get('status') == 'success':
                importance = 0.7
            if task_input.get('priority') == 'high':
                importance = 0.9
            
            # 确定标签
            tags = []
            if hasattr(task_type, 'value'):
                tags.append(task_type.value)
            else:
                tags.append(str(task_type))
            
            # 添加内容概览标签，允许跨分区共享
            task_text = str(task_input.get('input', ''))
            if any(keyword in task_text.lower() for keyword in ['总结', '摘要', '文档', '概述']):
                tags.append('content_overview')
            
            await context_partition_manager.add_context_entry(
                partition_type=partition_type,
                content=context_content,
                importance=importance,
                tags=tags,
                user_id=task_input.get('user_id')
            )
            self.logger.debug(f"已将任务 {task_id} 上下文添加到分区 {partition_type.value}")
        except Exception as e:
            self.logger.error(f"添加上下文到分区失败: {str(e)}")
            
    async def _execute_single_sub_task_async(self, sub_task: Dict[str, Any], task_id: str, task_embedding: List[float]) -> Dict[str, Any]:
        """异步执行单个子任务"""
        sub_task_id = sub_task["sub_task_id"]
        agent_type = sub_task["agent_type"]
        self.logger.info(f"异步执行子任务: {sub_task_id} | agent_type={agent_type}")

        # 更新子任务状态为执行中
        task_manager = TaskStateManager(task_id=task_id, redis_client=self.redis_client)
        task_manager.update_state(
            state=TaskState.RUNNING.value,
            agent_type=agent_type,
            input_params=sub_task
        )

        try:
            # 使用MATPO算法选择最优Agent实例
            instance = self._select_agent_by_matpo(task_embedding, agent_type)
            if not instance:
                # 回退到负载均衡选择
                instance = self._select_best_agent_by_load(agent_type)
                if not instance:
                    raise AgentNotFoundError(
                        message=f"No available agent instance for {agent_type}",
                        context={"agent_type": agent_type}
                    )
            
            instance["agent_type"] = agent_type  # 补充Agent类型，用于更新指标

            # 异步调用Agent实例执行子任务
            result = await self._call_agent_instance_async(instance, sub_task)

            if result["success"]:
                # 子任务成功：更新状态和Agent能力
                task_manager.update_state(
                    state=TaskState.SUCCESS.value,
                    agent_instance_id=instance["instance_id"],
                    output_result={"sub_task_id": sub_task_id, "result": result["data"]}
                )
                
                # 计算奖励并更新能力
                reward = 1.0  # 成功奖励
                self._update_agent_capabilities(instance["instance_id"], True, task_embedding, reward)
                
                return result
            else:
                # 子任务失败：触发重试
                self.logger.error(f"子任务{sub_task_id}执行失败，触发重试: {result['error_msg']}")
                retry_success, retry_result = await self._retry_task_async(task_id, agent_type, sub_task, task_embedding)
                if retry_success:
                    return retry_result
                else:
                    raise AgentCallError(message=f"Sub task {sub_task_id} failed after retries",
                                         context={"error_msg": retry_result["error_msg"]})
        except Exception as e:
            self.logger.error(f"子任务{sub_task_id}执行异常: {str(e)}")
            raise

    def _integrate_results(self, sub_task_results: List[Dict[str, Any]], original_input: str) -> str:
        """整合子任务结果（生成最终响应）"""
        if len(sub_task_results) == 1:
            # 单任务：直接返回结果
            return sub_task_results[0]["result"]
        else:
            # 多任务：整合所有子任务结果
            integrated = f"针对查询「{original_input}」的综合结果：\n"
            for idx, result in enumerate(sub_task_results, 1):
                integrated += f"{idx}. {result['result']}\n"
            return integrated.strip()

    def export_matpo_model(self) -> Dict[str, Any]:
        """导出MATPO模型参数（用于持久化和共享）"""
        model_data = {
            "agent_capabilities": {}
        }
        
        for instance_id, capability in self.agent_capabilities.items():
            model_data["agent_capabilities"][instance_id] = {
                "agent_type": capability.agent_type,
                "capability_vector": capability.capability_vector,
                "success_rate": capability.success_rate,
                "response_time": capability.response_time
            }
        
        return model_data

    def import_matpo_model(self, model_data: Dict[str, Any]) -> None:
        """导入MATPO模型参数"""
        if "agent_capabilities" in model_data:
            with self.cache_lock:
                for instance_id, cap_data in model_data["agent_capabilities"].items():
                    # 只导入存在的实例
                    if instance_id in self.agent_capabilities:
                        current_cap = self.agent_capabilities[instance_id]
                        # 平滑更新能力向量
                        current_cap.capability_vector = [
                            0.5 * current + 0.5 * imported
                            for current, imported in zip(current_cap.capability_vector, cap_data["capability_vector"])
                        ]
                        current_cap.success_rate = 0.5 * current_cap.success_rate + 0.5 * cap_data["success_rate"]
                        current_cap.response_time = 0.5 * current_cap.response_time + 0.5 * cap_data["response_time"]
                        
                        self.logger.debug(f"导入实例{instance_id}的MATPO模型参数")


# 创建全局MATPO协调器实例
matpo_orchestrator = None


def get_matpo_orchestrator() -> MATPOOrchestratorAgent:
    """获取MATPO协调器实例（单例模式）"""
    global matpo_orchestrator
    if matpo_orchestrator is None:
        matpo_orchestrator = MATPOOrchestratorAgent()
    return matpo_orchestrator


# ===================== 异步使用示例 =====================
async def main():
    """异步主函数示例"""
    import sys
    sys.path.append("../../")  # 添加项目根目录到Python路径

    # 初始化MATPO协调器
    orchestrator = get_matpo_orchestrator()
    
    # 等待1秒，让订阅线程获取Agent实例缓存
    await asyncio.sleep(1)

    # 示例1：执行单Agent任务（文本问答）
    print("=== 执行异步单Agent任务（文本问答）===")
    single_task_input = {
        "task_id": "matpo_single_task_001",
        "task_type": "single",
        "agent_type": AgentType.TEXT_QA.value,
        "input": "什么是MATPO多智能体训练方法？"
    }
    single_result = await orchestrator.execute_task_async(single_task_input)
    print(json.dumps(single_result, ensure_ascii=False, indent=2))

    # 示例2：执行多Agent协作任务（并发执行）
    print("\n=== 执行异步多Agent协作任务（并发）===")
    multi_task_input = {
        "task_id": "matpo_multi_task_001",
        "task_type": "multi",
        "input": "请分析MATPO算法在多智能体系统中的应用"
    }
    multi_result = await orchestrator.execute_task_async(multi_task_input)
    print(json.dumps(multi_result, ensure_ascii=False, indent=2))

    # 导出模型示例
    model_data = orchestrator.export_matpo_model()
    print("\n=== MATPO模型参数 ===")
    print(f"已学习的智能体能力数: {len(model_data['agent_capabilities'])}")


if __name__ == "__main__":
    # 运行异步示例
    asyncio.run(main())